from tensorflow.python.keras.callbacks import Callback
from typing import Union, List, Any


class TrainParam(object):
    def __init__(self,
                 batch_size: Union[None, int] = None,
                 epochs: int = 1,
                 verbose: int = 1,
                 callbacks: List[Callback] = None,
                 validation_split: float = 0.,
                 validation_data: Any = None,
                 shuffle: bool = True,
                 class_weight: Any = None,
                 sample_weight: Any = None,
                 initial_epoch: int = 0,
                 steps_per_epoch: Union[None, int] = None,
                 validation_steps: Any = None,
                 max_queue_size: int = 10,
                 workers: int = 1,
                 use_multiprocessing: bool = False) -> None:
        self.use_multiprocessing = use_multiprocessing
        self.workers = workers
        self.max_queue_size = max_queue_size
        self.validation_steps = validation_steps
        self.steps_per_epoch = steps_per_epoch
        self.initial_epoch = initial_epoch
        self.sample_weight = sample_weight
        self.class_weight = class_weight
        self.shuffle = shuffle
        self.validation_data = validation_data
        self.validation_split = validation_split
        self.callbacks = callbacks
        self.verbose = verbose
        self.epochs = epochs
        self.batch_size = batch_size
